{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing parse_assistant_message Function\n",
    "\n",
    "This notebook demonstrates the functionality of the `parse_assistant_message` function from the `coder.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2024-12-06 21:31:20.746\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mbyzerllm.utils.connect_ray\u001b[0m:\u001b[36mconnect_cluster\u001b[0m:\u001b[36m48\u001b[0m - \u001b[1mJDK 21 will be used (/Users/allwefantasy/.auto-coder/jdk-21.0.2.jdk/Contents/Home)...\u001b[0m\n",
      "2024-12-06 21:31:20,788\tINFO worker.py:1564 -- Connecting to existing Ray cluster at address: 127.0.0.1:6379...\n",
      "2024-12-06 21:31:20,803\tINFO worker.py:1740 -- Connected to Ray cluster. View the dashboard at \u001b[1m\u001b[32m127.0.0.1:8265 \u001b[39m\u001b[22m\n"
     ]
    }
   ],
   "source": [
    "from autocoder.agent.coder import Coder\n",
    "import byzerllm\n",
    "\n",
    "# Create a mock instance just to access the parse_assistant_message function\n",
    "llm = byzerllm.ByzerLLM.from_default_model(model=\"deepseek_chat\")\n",
    "coder = Coder(llm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Case 1: Basic Text Content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Case 1 Result:\n",
      "[{'type': 'text', 'content': 'Hello, this is a simple text message', 'partial': True}]\n"
     ]
    }
   ],
   "source": [
    "message = \"Hello, this is a simple text message\"\n",
    "result = coder.parse_assistant_message(message)\n",
    "print(\"Test Case 1 Result:\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Case 2: Single Tool Use (execute_command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Case 2 Result:\n",
      "[{'type': 'text', 'content': \"I'll execute the ls command:\", 'partial': False}, {'type': 'tool_use', 'name': 'execute_command', 'params': {'command': 'ls -l'}, 'partial': False}, {'type': 'text', 'content': '', 'partial': True}]\n"
     ]
    }
   ],
   "source": [
    "message = \"\"\"\n",
    "I'll execute the ls command:\n",
    "<execute_command>\n",
    "<command>ls -l</command>\n",
    "</execute_command>\n",
    "\"\"\"\n",
    "result = coder.parse_assistant_message(message)\n",
    "print(\"Test Case 2 Result:\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Case 3: Multiple Tools with Text In Between"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Case 3 Result:\n",
      "[{'type': 'text', 'content': \"First, let's read a file:\", 'partial': False}, {'type': 'tool_use', 'name': 'read_file', 'params': {'path': 'test.txt'}, 'partial': False}, {'type': 'text', 'content': \"Now, let's write to another file:\", 'partial': False}, {'type': 'tool_use', 'name': 'write_to_file', 'params': {'path': 'output.txt', 'content': 'Hello World!\\nThis is a test.'}, 'partial': False}, {'type': 'text', 'content': '', 'partial': True}]\n"
     ]
    }
   ],
   "source": [
    "message = \"\"\"\n",
    "First, let's read a file:\n",
    "<read_file>\n",
    "<path>test.txt</path>\n",
    "</read_file>\n",
    "\n",
    "Now, let's write to another file:\n",
    "\n",
    "<write_to_file>\n",
    "<path>output.txt</path>\n",
    "<content>\n",
    "Hello World!\n",
    "This is a test.\n",
    "</content>\n",
    "</write_to_file>\n",
    "\"\"\"\n",
    "result = coder.parse_assistant_message(message)\n",
    "print(\"Test Case 3 Result:\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Case 4: Tool With Multiple Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Case 4 Result:\n",
      "[{'type': 'text', 'content': '', 'partial': False}, {'type': 'tool_use', 'name': 'search_files', 'params': {'path': 'src', 'regex': 'def\\\\s+\\\\w+', 'file_pattern': '*.py'}, 'partial': False}, {'type': 'text', 'content': '', 'partial': True}]\n"
     ]
    }
   ],
   "source": [
    "message = \"\"\"\n",
    "<search_files>\n",
    "<path>src</path>\n",
    "<regex>def\\s+\\w+</regex>\n",
    "<file_pattern>*.py</file_pattern>\n",
    "</search_files>\n",
    "\"\"\"\n",
    "result = coder.parse_assistant_message(message)\n",
    "print(\"Test Case 4 Result:\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Case 5: Attempt Completion Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Case 5 Result:\n",
      "[{'type': 'text', 'content': '', 'partial': False}, {'type': 'tool_use', 'name': 'attempt_completion', 'params': {'result': 'Task completed successfully! Here are the results:\\n1. Files were processed\\n2. Data was analyzed', 'command': 'cat output.txt'}, 'partial': False}, {'type': 'text', 'content': '', 'partial': True}]\n"
     ]
    }
   ],
   "source": [
    "message = \"\"\"\n",
    "<attempt_completion>\n",
    "<result>\n",
    "Task completed successfully! Here are the results:\n",
    "1. Files were processed\n",
    "2. Data was analyzed\n",
    "</result>\n",
    "<command>cat output.txt</command>\n",
    "</attempt_completion>\n",
    "\"\"\"\n",
    "result = coder.parse_assistant_message(message)\n",
    "print(\"Test Case 5 Result:\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Case 6: Partial Tool Use (Simulating Streaming)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Case 6 Result (Partial Tool Use):\n",
      "[{'type': 'text', 'content': '', 'partial': False}, {'type': 'tool_use', 'name': 'execute_command', 'params': {'command': 'ls'}, 'partial': True}]\n"
     ]
    }
   ],
   "source": [
    "message = \"\"\"\n",
    "<execute_command>\n",
    "<command>ls\n",
    "\"\"\"\n",
    "result = coder.parse_assistant_message(message)\n",
    "print(\"Test Case 6 Result (Partial Tool Use):\")\n",
    "print(result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
